import re
import time
import paramiko
import logging
from fabric import Connection

# 远程模式


class SSH:
    def __init__(self, host, username, password, port=22):
        self.ssh = paramiko.SSHClient()
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.ssh.connect(hostname=host, username=username, password=password, port=port)
        self.channel = self.ssh.invoke_shell()
        self.sftp = self.ssh.open_sftp()
        time.sleep(0.5)

    def cmd(self, cmd, keywords=r"\[.*\]\#\s*$", times=60):
        """执行命令"""
        count = 0
        rrecv = ""
        try:
            self.channel.send(cmd + "\n")
            while count < times:
                time.sleep(0.5)
                buf = self.channel.recv(1024 * 1024).decode(
                    encoding="utf-8", errors="ignore"
                )
                # 处理特殊字符和 ANSI 控制字符
                buf = re.sub(r"\x1b\[([0-9]{1,2}(;[0-9]{1,2})?)?[m|K]?", "", buf)
                # 处理Last login
                buf = re.sub(r"Last login:.*from\s\d+\.\d+\.\d+\.\d+\r\r\n", "", buf)
                count += 1
                logging.info(buf)
                rrecv += buf
                if re.search(keywords, buf):
                    break
        except Exception as ex:
            logging.error("Error:Command Error!")
            rrecv += str(ex)
        return rrecv

    def close(self):
        self.sftp.close()
        self.channel.close()
        self.ssh.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class SSH2:
    def __init__(self, host, username, password) -> None:
        # 初始化连接
        self.ssh = Connection(
            host=host, user=username, connect_kwargs={"password": password}
        )

    def cmd(self, cmd):
        # 执行远程脚本
        try:
            result = self.ssh.run(cmd, hide=True, encoding="utf-8")
        except Exception as e:
            result = None
            print(f"执行命令 {cmd} 出错，错误信息：{e}")
        return result

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.ssh.close()

    def __enter__(self):
        return self
